#include <omp.h>

#include <cuda_runtime.h>

#include <stdio.h>

#define min(x, y) ( ((x) < (y)) ? (x) : (y) )

void Kernel_Wrapper(float *, int, int);

int main(int argc, char * argv[])
{
	int device_count = 0;
	int N = 32;

	cudaGetDeviceCount(&device_count);
	omp_set_num_threads(device_count);

	float *shared_mem = new float[device_count];

#pragma omp parallel
	{		
		unsigned int cpu_thread_id = omp_get_thread_num();
		unsigned int num_cpu_threads = omp_get_num_threads();

		cudaSetDevice(cpu_thread_id);

		fprintf(stdout, "Rank %d size %d device %d / %d\n", cpu_thread_id, num_cpu_threads, device_count);

		float * dData = NULL;
		float * hData = NULL;

		cudaMallocHost((void **) &hData, N * sizeof(float));
		cudaMalloc((void **) &dData, N * sizeof(float));

		Kernel_Wrapper(dData, cpu_thread_id, N);

		cudaMemcpy(hData, dData, N * sizeof(float), cudaMemcpyDeviceToHost);

		fprintf(stdout, "0 : %f | 1 : %f | 2 : %f | 3 : %f\n", hData[0], hData[1], hData[2], hData[3]);

		shared_mem[(cpu_thread_id + 1) % num_cpu_threads] = hData[3];

#pragma omp barrier

		fprintf(stdout, "Thread %d recv %f\n", cpu_thread_id, shared_mem[cpu_thread_id]);

		cudaFree(dData);
		cudaFreeHost(hData);
	}

	delete [] shared_mem;

	return 0;
}